package ai

import (
	"app/pkg/log"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
)

type SparkAi struct {
	baseUrl   string
	appId     string
	apiKey    string
	apiSecret string
}

func NewSparkAi(url, appid, apikey, apiSecret string) *SparkAi {
	return &SparkAi{url, appid, apikey, apiSecret}
}

func (ai SparkAi) Chat(ctx *gin.Context, msg []*ChatMessage) (string, error) {
	d := websocket.Dialer{
		HandshakeTimeout: 10 * time.Second,
	}
	//TODO 消息的组装和长度控制
	messages := []ConversationMessage{}

	var resp *http.Response
	target := ai.assembleAuthUrl1(ai.baseUrl, ai.apiKey, ai.apiSecret)
	//握手并建立websocket 连接
	conn, resp, err := d.Dial(target, make(http.Header))
	if err != nil {
		return "", errors.New("无法建立连接")
	}
	if resp.StatusCode != 101 {
		return "", errors.New("!101")
	}
	err = conn.WriteJSON(map[string]interface{}{
		"header": map[string]interface{}{
			"app_id": ai.appId,
		},
		"parameter": map[string]interface{}{
			"chat": map[string]interface{}{
				"domain":      "general",
				"temperature": float64(0.8),
				"max_tokens":  1024,
				//"top_k":       int64(6),
				//"max_tokens":  int64(2048),
				//"auditing": "default",
			},
		},
		"payload": map[string]interface{}{
			"message": map[string]interface{}{
				"text": messages,
			},
		},
	})
	if err != nil {
		return "", errors.New("消息发送失败!" + err.Error())
	}
	return ai.StreamResponse(ctx, conn), nil
}

func (ai SparkAi) StreamResponse(ctx *gin.Context, conn *websocket.Conn) string {
	w := ctx.Writer
	var builder strings.Builder
	//设置stream响应头
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.WriteHeader(http.StatusOK)

	var index int
	var ticker = time.NewTicker(50 * time.Millisecond)
	var log = log.Logger()
	defer ticker.Stop()
loop:
	for {
		select {
		case <-ctx.Request.Context().Done():
			break loop
		case <-ticker.C:
			_, msg, er := conn.ReadMessage()
			if er != nil {
				log.Errorf(ctx, "读取消息报错：%s", er)
				break loop
			}
			var resp XunFeiStreamResponse
			if err := json.Unmarshal(msg, &resp); err != nil {
				continue
			}
			if resp.Header.Code != 0 {
				//错误码，0表示正常，非0表示出错；详细释义可在接口说明文档最后的错误码说明了解
				log.Errorf(ctx, "会话[%s]报错：%s", resp.Header.Sid, resp.Header.Code)
				break loop
			}
			if resp.Payload.Choices.Seq > index {
				index = resp.Payload.Choices.Seq
				content := resp.Payload.Choices.Text[0].Content
				builder.WriteString(content)
				_, _ = w.WriteString(content)
				w.(http.Flusher).Flush()
			}
			if resp.Header.Status == 2 {
				break loop
			}
		}
	}
	return builder.String()
}

// 创建鉴权url  apikey 即 hmac username
func (ai SparkAi) assembleAuthUrl1(hosturl string, apiKey, apiSecret string) string {
	ul, err := url.Parse(hosturl)
	if err != nil {
		fmt.Println(err)
	}
	//签名时间
	date := time.Now().UTC().Format(time.RFC1123)
	//date = "Tue, 28 May 2019 09:10:42 MST"
	//参与签名的字段 host ,date, request-line
	signString := []string{"host: " + ul.Host, "date: " + date, "GET " + ul.Path + " HTTP/1.1"}
	//拼接签名字符串
	sgin := strings.Join(signString, "\n")
	// fmt.Println(sgin)
	//签名结果
	sha := HmacWithShaTobase64("hmac-sha256", sgin, apiSecret)
	// fmt.Println(sha)
	//构建请求参数 此时不需要urlencoding
	authUrl := fmt.Sprintf("hmac username=\"%s\", algorithm=\"%s\", headers=\"%s\", signature=\"%s\"", apiKey,
		"hmac-sha256", "host date request-line", sha)
	//将请求参数使用base64编码
	authorization := base64.StdEncoding.EncodeToString([]byte(authUrl))

	v := url.Values{}
	v.Add("host", ul.Host)
	v.Add("date", date)
	v.Add("authorization", authorization)
	//将编码后的字符串url encode后添加到url后面
	callurl := hosturl + "?" + v.Encode()
	return callurl
}

type XunFeiStreamRequest struct {
}
type XunFeiStreamResponse struct {
	Header struct {
		Code    int    `json:"code"`
		Message string `json:"message"`
		Sid     string `json:"sid"`
		Status  int    `json:"status"`
	} `json:"header"`
	Payload struct {
		Choices struct {
			Status int `json:"status"`
			Seq    int `json:"seq"`
			Text   []struct {
				Index   int    `json:"index"`
				Role    string `json:"role"`
				Content string `json:"content"`
			} `json:"text"`
		} `json:"choices"`
		Usage struct {
			Text struct {
				QuestionTokens   int `json:"question_tokens"`
				PromptTokens     int `json:"prompt_tokens"`
				CompletionTokens int `json:"completion_tokens"`
				TotalTokens      int `json:"total_tokens"`
			} `json:"text"`
		} `json:"usage"`
	} `json:"payload"`
}

func HmacWithShaTobase64(algorithm, data, key string) string {
	mac := hmac.New(sha256.New, []byte(key))
	mac.Write([]byte(data))
	encodeData := mac.Sum(nil)
	return base64.StdEncoding.EncodeToString(encodeData)
}
